#  -*- coding: UTF-8 -*-

# MindPlus
# Python
import os
import cv2
import numpy as np
from threading import Lock
from PIL import Image, ImageDraw, ImageFont



class CameraManager:
    _instance = None
    _lock = Lock()

    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._init_camera()
        return cls._instance

    def _init_camera(self):
        self.cap = cv2.VideoCapture(0)
        if not self.cap.isOpened():
            raise RuntimeError("摄像头初始化失败")
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 800)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

    def get_frame(self):
        with self._lock:
            ret, frame = self.cap.read()
            if not ret:
                raise RuntimeError("无法读取摄像头帧")
            return frame

    def release(self):
        if self.cap.isOpened():
            self.cap.release()

    def __del__(self):
        self.release()

    def is_opened(self):
        return hasattr(self, 'cap') and self.cap.isOpened()

    def restart_camera(self):
        with self._lock:
            if hasattr(self, 'cap') and self.cap.isOpened():
                self.cap.release()
            self._init_camera()

# 全局摄像头管理器实例
if 'camera_manager' not in globals():
    camera_manager = CameraManager()


# 初始化人脸检测模型
face_net = cv2.dnn.readNetFromCaffe(
    "wenjian/deploy.prototxt",
    "wenjian/res10_300x300_ssd_iter_140000.caffemodel"
)

# 初始化识别器
recognizer = cv2.face.LBPHFaceRecognizer_create()
if os.path.exists("wenjian/trainer.yml"):
    recognizer.read("wenjian/trainer.yml")

# 加载字体
font = ImageFont.truetype('wenjian/simhei.ttf', 20)

# 初始化用户信息
names = []
roles = []
if os.path.exists("name.txt"):
    with open("name.txt", 'r', encoding='utf-8') as f:
        lines = [line.strip() for line in f.readlines()]
        names = lines[:5]
        roles = lines[5:]
elif os.path.exists("label_mapping.txt"):
    with open("label_mapping.txt", 'r', encoding='utf-8') as f:
        for line in f:
            label, name_role = line.strip().split(':', 1)
            if '_' in name_role:
                name, role = name_role.split('_', 1)
            else:
                name, role = name_role, ""
            names.append(name)
            roles.append(role)


# 创建数据目录
data_dir = 'data'
if not os.path.exists(data_dir):
    os.makedirs(data_dir)

# 清空当前用户的数据
user_dir = os.path.join(data_dir, f'"张三"_"非遗传承人"')
if os.path.exists(user_dir):
    for root, dirs, files in os.walk(user_dir, topdown=False):
        for file in files:
            os.remove(os.path.join(root, file))
        for dir in dirs:
            os.rmdir(os.path.join(root, dir))
    os.rmdir(user_dir)
os.makedirs(user_dir)

# 确保摄像头已初始化
if not camera_manager.is_opened():
    camera_manager.restart_camera()

count = 0
print(f"开始采集 "张三" 的人脸数据...")

while count < 15:
    try:
        frame = camera_manager.get_frame()
        frame = cv2.rotate(frame, cv2.ROTATE_180)
        frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(frame_pil)

        # 人脸检测
        (h, w) = frame.shape[:2]
        blob = cv2.dnn.blobFromImage(cv2.resize(frame, (300, 300)), 1.0,
                                    (300, 300), (104.0, 177.0, 123.0))
        face_net.setInput(blob)
        detections = face_net.forward()
        faces = []
        for i in range(0, detections.shape[2]):
            confidence = detections[0, 0, i, 2]
            if confidence > 0.5:
                box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
                (startX, startY, endX, endY) = box.astype("int")
                faces.append((startX, startY, endX - startX, endY - startY))

        if len(faces) == 1:
            x, y, w, h = faces[0]
            draw.rectangle([x, y, x+w, y+h], outline=(0, 255, 0), width=2)
            draw.text((x, y-25), f"检测到人脸，按空格拍摄 ({count}/15)",
                     font=font, fill=(0, 255, 0))
        else:
            draw.text((10, 20), "未检测到人脸", font=font, fill=(255, 0, 0))

        draw.text((10, 40), f"当前用户: "张三" ("非遗传承人")",
                 font=font, fill=(0, 255, 0))
        draw.text((10, 60), "ESC退出 | 空格拍摄",
                 font=font, fill=(0, 255, 0))

        cv2.imshow('Face Capture', cv2.cvtColor(np.array(frame_pil), cv2.COLOR_RGB2BGR))

        key = cv2.waitKey(30)
        if key == 32 and len(faces) == 1:  # 空格键拍照
            x, y, w, h = faces[0]
            face_roi = frame[y:y+h, x:x+w]
            gray = cv2.cvtColor(face_roi, cv2.COLOR_BGR2GRAY)
            cv2.imwrite(os.path.join(user_dir, f'"张三"_"非遗传承人"_{count}.jpg'), gray)
            count += 1
            print(f"已保存第 {count} 张照片")
        elif key == 27:  # ESC退出
            break
    except Exception as e:
        print(f"摄像头错误: {str(e)}")
        camera_manager.restart_camera()

print(f""张三" 的人脸数据采集完成!")
cv2.destroyAllWindows()

# 更新用户信息
if not os.path.exists("name.txt"):
    with open("name.txt", 'w', encoding='utf-8') as f:
        f.write(f""张三"\n"非遗传承人"")


# 创建数据目录
data_dir = 'data'
if not os.path.exists(data_dir):
    os.makedirs(data_dir)

# 清空当前用户的数据
user_dir = os.path.join(data_dir, f'"李四"_"传统手艺人"')
if os.path.exists(user_dir):
    for root, dirs, files in os.walk(user_dir, topdown=False):
        for file in files:
            os.remove(os.path.join(root, file))
        for dir in dirs:
            os.rmdir(os.path.join(root, dir))
    os.rmdir(user_dir)
os.makedirs(user_dir)

# 确保摄像头已初始化
if not camera_manager.is_opened():
    camera_manager.restart_camera()

count = 0
print(f"开始采集 "李四" 的人脸数据...")

while count < 15:
    try:
        frame = camera_manager.get_frame()
        frame = cv2.rotate(frame, cv2.ROTATE_180)
        frame_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(frame_pil)

        # 人脸检测
        (h, w) = frame.shape[:2]
        blob = cv2.dnn.blobFromImage(cv2.resize(frame, (300, 300)), 1.0,
                                    (300, 300), (104.0, 177.0, 123.0))
        face_net.setInput(blob)
        detections = face_net.forward()
        faces = []
        for i in range(0, detections.shape[2]):
            confidence = detections[0, 0, i, 2]
            if confidence > 0.5:
                box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
                (startX, startY, endX, endY) = box.astype("int")
                faces.append((startX, startY, endX - startX, endY - startY))

        if len(faces) == 1:
            x, y, w, h = faces[0]
            draw.rectangle([x, y, x+w, y+h], outline=(0, 255, 0), width=2)
            draw.text((x, y-25), f"检测到人脸，按空格拍摄 ({count}/15)",
                     font=font, fill=(0, 255, 0))
        else:
            draw.text((10, 20), "未检测到人脸", font=font, fill=(255, 0, 0))

        draw.text((10, 40), f"当前用户: "李四" ("传统手艺人")",
                 font=font, fill=(0, 255, 0))
        draw.text((10, 60), "ESC退出 | 空格拍摄",
                 font=font, fill=(0, 255, 0))

        cv2.imshow('Face Capture', cv2.cvtColor(np.array(frame_pil), cv2.COLOR_RGB2BGR))

        key = cv2.waitKey(30)
        if key == 32 and len(faces) == 1:  # 空格键拍照
            x, y, w, h = faces[0]
            face_roi = frame[y:y+h, x:x+w]
            gray = cv2.cvtColor(face_roi, cv2.COLOR_BGR2GRAY)
            cv2.imwrite(os.path.join(user_dir, f'"李四"_"传统手艺人"_{count}.jpg'), gray)
            count += 1
            print(f"已保存第 {count} 张照片")
        elif key == 27:  # ESC退出
            break
    except Exception as e:
        print(f"摄像头错误: {str(e)}")
        camera_manager.restart_camera()

print(f""李四" 的人脸数据采集完成!")
cv2.destroyAllWindows()

# 更新用户信息
if not os.path.exists("name.txt"):
    with open("name.txt", 'w', encoding='utf-8') as f:
        f.write(f""李四"\n"传统手艺人"")


# 确保摄像头已初始化
if not camera_manager.is_opened():
    camera_manager.restart_camera()

cv2.namedWindow('Face Recognition', cv2.WINDOW_NORMAL)
cv2.resizeWindow('Face Recognition', 800, 480)
cv2.moveWindow('Face Recognition', 0, 0)

while True:
    try:
        frame = camera_manager.get_frame()
        frame = cv2.rotate(frame, cv2.ROTATE_180)
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frame_pil = Image.fromarray(frame_rgb)
        draw = ImageDraw.Draw(frame_pil)

        # 人脸检测
        (h, w) = frame.shape[:2]
        blob = cv2.dnn.blobFromImage(cv2.resize(frame, (300, 300)), 1.0,
                                    (300, 300), (104.0, 177.0, 123.0))
        face_net.setInput(blob)
        detections = face_net.forward()
        faces = []
        for i in range(0, detections.shape[2]):
            confidence = detections[0, 0, i, 2]
            if confidence > 0.5:
                box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
                (startX, startY, endX, endY) = box.astype("int")
                faces.append((startX, startY, endX - startX, endY - startY))

        if len(faces) == 1:
            x, y, w, h = faces[0]
            draw.rectangle([x, y, x + w, y + h], outline=(0, 255, 0), width=2)

            # 识别当前人脸
            face_roi = frame[y:y + h, x:x + w]
            gray = cv2.cvtColor(face_roi, cv2.COLOR_BGR2GRAY)
            id_, confidence = recognizer.predict(gray)

            # 显示识别结果
            if confidence < 70:  # 置信度阈值
                name = names[id_] if id_ < len(names) else "未知"
                role = roles[id_] if id_ < len(roles) else ""
                display_text = f"{name} {role}"
                draw.text((x + w + 10, y), display_text, font=font, fill=(0, 255, 0))

                # 保存识别结果
                with open("wenjian/face_result.txt", "w", encoding="utf-8") as f:
                    f.write(f"{name}\n{role}")
            else:
                draw.text((x + w + 10, y), "未识别到用户", font=font, fill=(255, 0, 0))

        cv2.imshow('Face Recognition', cv2.cvtColor(np.array(frame_pil), cv2.COLOR_RGB2BGR))

        key = cv2.waitKey(30)
        if key == 27:  # ESC退出
            break
    except Exception as e:
        print(f"摄像头错误: {str(e)}")
        camera_manager.restart_camera()

cv2.destroyAllWindows()


print("开始训练人脸识别模型...")

# 获取训练数据
image_paths = []
labels = []
label_dict = {}

# 遍历data目录下的所有用户文件夹
for user_dir in os.listdir('data'):
    user_path = os.path.join('data', user_dir)
    if not os.path.isdir(user_path):
        continue

    # 直接使用文件夹名作为用户ID
    user_id = user_dir
    user_path = os.path.join('data', user_dir)

    # 遍历用户文件夹中的所有图片
    for img_name in os.listdir(user_path):
        if not img_name.endswith('.jpg'):
            continue

        img_path = os.path.join(user_path, img_name)
        image_paths.append((img_path, user_id))

# 读取所有图片并创建标签
faces = []
labels = []
current_label = 0

for img_path, user_id in image_paths:
    # 读取灰度图像
    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
    if img is None:
        continue

    # 为每个用户分配一个数字标签
    if user_id not in label_dict:
        label_dict[user_id] = current_label
        current_label += 1

    faces.append(img)
    labels.append(label_dict[user_id])

if not faces:
    print("错误：没有找到训练数据！")
else:
    # 训练模型
    recognizer.train(faces, np.array(labels))

    # 保存模型
    recognizer.save("wenjian/trainer.yml")
    print("模型训练完成，已保存到 wenjian/trainer.yml")

    # 保存标签映射关系
    with open('label_mapping.txt', 'w', encoding='utf-8') as f:
        for name, label in label_dict.items():
            f.write(f"{label}:{name}\n")

    print("训练成功！")

print((text := open("wenjian/face_result.txt", "r", encoding="utf-8").read().split("\n")[1].strip()) and text or "")
print((text := open("wenjian/face_result.txt", "r", encoding="utf-8").read().split("\n")[1].strip()) and text or "")
